#include "Image.H"
#include "IO.H"
#include "Util.H"

#include "DataSet.h"
#include "kNNClassifier.h"
#include "NearestMeanClassifier.h"
#include "Sample.h"
#include "TextureClassification.h"

#include <iostream>
#include <sstream>
#include <string>
#include <vector>

using namespace pip;
using namespace std;

//-----------------------------------------------//

// Little helper, pushes file names in a vector
void getFiles(const string& base, vector<string>& v)
{
  for (int c = 0; c < 10; ++c) {
    for (int i = 0; i < 8; ++i) {
      string file;
      stringstream str;
      str << base << c << "-" << i << ".pip";
      str >> file;
      v.push_back(file);
    }
  }
}

//-----------------------------------------------//

// Helper for experiment 1
void getFeatureLabels(const vector<string>& files,
                      const int& bins,
                      DataSet& ds
                      )
{
  const unsigned int nrFiles(files.size());
  for (unsigned int i = 0; i < nrFiles; ++i) {
    vector<float> v;
    Image<float> in;
    if (!importFile(in, files[i])) {
      error("", "can't read file");
    }
    histogram(in, v, bins);
    Sample s(v, i / 8);
    ds.push_back(s);
  }
}

//-----------------------------------------------//

// Helper for experiment 2
void getLawFeatureLabels(const vector<string>& files,
                         const string& filters,
                         const int& bins, 
                         DataSet& ds
                         )
{
  const unsigned int nrFiles(files.size());
  const unsigned int nrFilters(filters.size());
  for (unsigned int i = 0; i < nrFiles; ++i) {
    vector<float> v;
    Image<float> in, out;
    if (!importFile(in, files[i])) {
      error("", "can't read file");
    }

    // Compute the desired law3 filters
    for (unsigned int j = 0; j < nrFilters; j += 2) {
      vector<float> temp;
      Law3(in, out, filters.substr(j, 2)); 
      histogram(out, temp, bins);
      v.insert(v.end(), temp.begin(), temp.end());
    }

    Sample s(v, i / 8);
    ds.push_back(s);
  }
}

//------------------------------------------------//

void experiment(const DataSet& trainingds,
                const DataSet& testds,
                const vector<string> &names,
                const bool& nm, 
                const bool& knn, 
                const int& k = 1
                )
{
  if (trainingds.samples().size() <= 0) {
    cerr << "Not enough samples" << endl;
    return;
  }

  cout << "Classification experiment with " << trainingds.samples().size() <<
    " training samples, " << trainingds.sample(0).input().size() << " features, " <<
    testds.samples().size() << " test samples." << endl;

  float accuracy;

  Classifier* c;
  if (nm) {
    c = new NearestMeanClassifier();
  } else if (knn) {
    c = new kNNClassifier(k);
  }
  c->train(trainingds);
  c->experiment(testds, accuracy, names, true);

  delete c;

  cout << "Accuracy: " << accuracy << endl;
}

//---------------------------------------------------//

int main(int argc, char *argv[])
{
  const string program = argv[0];
  vector<string> args(argv, argv + argc);
  if (argc < 2) {
    error(program, "Usage: texture\n Texture assignment\n");
  }

  unsigned int bins = 5;
  getScalarArg(args, "-bins", bins);

  bool his = false;
  getFlagArg(args, "-histogram", his);
  if (his) {
    Image<float> in;
    if (!importFile(in, argv[1])) {
      error(program, "can't read file");
    }

    vector<float> v;
    histogram(in, v, bins);

    for (unsigned int i = 0; i < bins; ++i) {
      cout << "Bin " << i << ": " << v[i] << endl;
    }
  }

  //flags for classifiers
  bool nm = false; // use nearest mean classifier
  getFlagArg(args, "-nm", nm);
  bool knn = false; // use k nearest neighbors classifier
  getFlagArg(args, "-knn", knn);

  // the k to use in knn classification
  int k = 1;
  getScalarArg( args, "-k", k);

  //variables and calls for experiments
  vector<string> trainfiles, testfiles;
  getFiles("train", trainfiles);
  getFiles("test", testfiles);
  DataSet trainingds, testds;

  bool exp1 = false;
  getFlagArg(args, "-exp1", exp1);
  if (exp1) {
    //get features, labels of training data
    getFeatureLabels(trainfiles, bins, trainingds);
    cout << "extracted features from training images" << endl;

    //idem for the test data
    getFeatureLabels(testfiles, bins, testds);
    cout << "extracted features from test images" << endl;

    //train & test the classifiers
    experiment(trainingds, testds, testfiles, nm, knn, k);
  }

  bool prep = false;
  getFlagArg(args, "-preprocess", prep);
  if (prep) {
    vector<string> files;
    getFiles("train", files);
    getFiles("test", files);

    const unsigned int nrFiles(files.size());
    for (unsigned int i = 0; i < nrFiles; ++i) {
      Image<float> in, out;
      if (!importFile(in, files[i])) {
        error(program,"can't read file");
      }

      preprocess(in,out);

      //NOTE: WILL OVERWRITE INPUT
      if (!exportFile(out, files[i])) {
        error(program, "can't write file");
      }
      cout << "processed " << files[i] << endl;
    }
  }

  bool law3 = false;
  getFlagArg(args, "-law3", law3);

  string kernel = "LL";
  getScalarArg(args, "-kernel", kernel);

  string outputFile = "law3.pip";
  getScalarArg(args, "-out", outputFile);

  if (law3) {
    if(outputFile == "law3.pip") { //no output filename was given
      cout << "No output file name specified with -out. Using law3.pip" << endl;
    }

    Image<float> in, out;
    if (!importFile(in, argv[1])) {
      error(program, "can't read file");
    }

    Law3(in, out, kernel);
    
    if (!exportFile(out, outputFile)) {
      error(program, "can't write file");
    }
  }

  bool exp2 = false;
  getFlagArg( args, "-exp2", exp2);
  if (exp2) {
    //get features, labels of training data
    getLawFeatureLabels(trainfiles, kernel, bins, trainingds);
    cout << "extracted features from training images with kernelstring " << kernel << endl;

    //idem for the test data
    getLawFeatureLabels(testfiles, kernel, bins, testds);
    cout << "extracted features from test images with kernelstring " << kernel << endl;

    //train & test the classifiers
    experiment(trainingds, testds, testfiles, nm, knn, k);
  }

  //system("pause");
  return 0;
}